"""
遗传算法模块 - 阶段2: 遗传算法优化 (核心代码模块)
实现定日镜场布局的遗传算法优化 - 修复版本
"""

import numpy as np
import random
import math
from data_preprocess import DataPreprocessor
from efficiency_calc import EfficiencyCalculator

class GeneticAlgorithm:
    def __init__(self, candidate_positions, solar_positions, dni_avg):
        """初始化遗传算法"""
        # 遗传算法参数 (标准参数)
        self.population_size = 100  # 标准种群大小
        self.max_generations = 200  # 标准迭代次数
        self.crossover_rate = 0.8   # 交叉概率
        self.mutation_rate = 0.1    # 标准变异概率
        self.selection_method = "tournament"  # 锦标赛选择
        
        # 问题参数
        self.candidate_positions = candidate_positions
        self.solar_positions = solar_positions
        self.dni_avg = dni_avg
        
        # 约束参数 - 标准: 100 ≤ N ≤ 2000
        self.min_mirrors = 100  # 标准最小镜数
        self.max_mirrors = 2000  # 标准最大镜数
        self.R_field = 350
        self.safety_distance = 11  # 标准: max(w,h) + 5m = max(6,6) + 5 = 11m
        
        # 场地面积 (用于目标函数计算)
        self.field_area = math.pi * self.R_field**2
        
        # 效率计算器
        self.efficiency_calc = EfficiencyCalculator()
        
        print("✅ 遗传算法初始化完成 (修复版本)")
        print(f"   种群大小: {self.population_size}")
        print(f"   最大迭代: {self.max_generations}")
        print(f"   安全距离: {self.safety_distance}m (已优化)")
        print(f"   候选位置: {len(self.candidate_positions)} 个")
    
    def create_individual(self):
        """创建个体 (染色体编码) - 修复版本"""
        # 修复: 动态调整定日镜数量范围
        max_possible_mirrors = len(self.candidate_positions)
        actual_min_mirrors = min(self.min_mirrors, max_possible_mirrors)
        actual_max_mirrors = min(self.max_mirrors, max_possible_mirrors)
        
        # 确保最小值不超过最大值
        if actual_min_mirrors > actual_max_mirrors:
            actual_min_mirrors = max(1, actual_max_mirrors // 2)
        
        # 随机选择定日镜数量
        N = random.randint(actual_min_mirrors, actual_max_mirrors)
        
        # 随机选择定日镜位置索引 (避免重复)
        mirror_indices = random.sample(range(len(self.candidate_positions)), N)
        
        # 随机选择吸收塔位置
        tower_x = random.uniform(-self.R_field, self.R_field)
        tower_y = random.uniform(-self.R_field, self.R_field)
        
        # 确保吸收塔在场地内
        while tower_x**2 + tower_y**2 > self.R_field**2:
            tower_x = random.uniform(-self.R_field, self.R_field)
            tower_y = random.uniform(-self.R_field, self.R_field)
        
        # 染色体编码: [N, index1, index2, ..., indexN, tower_x, tower_y]
        individual = [N] + mirror_indices + [tower_x, tower_y]
        
        return individual
    
    def initialize_population(self):
        """初始化种群"""
        population = []
        
        print("🔄 正在初始化种群...")
        
        for i in range(self.population_size):
            individual = self.create_individual()
            population.append(individual)
            
            if (i + 1) % 20 == 0:
                print(f"   已生成 {i + 1}/{self.population_size} 个个体")
        
        print("✅ 种群初始化完成")
        return population
    
    def decode_individual(self, individual):
        """解码个体 (提取定日镜位置和吸收塔位置) - 修复版本"""
        N = individual[0]
        mirror_indices = individual[1:N+1]
        tower_x, tower_y = individual[N+1], individual[N+2]
        
        # 修复: 检查索引有效性，过滤无效索引
        valid_indices = []
        for idx in mirror_indices:
            if 0 <= idx < len(self.candidate_positions):
                valid_indices.append(idx)
        
        # 如果有效索引不足，补充随机索引
        while len(valid_indices) < N and len(valid_indices) < len(self.candidate_positions):
            new_idx = random.randint(0, len(self.candidate_positions) - 1)
            if new_idx not in valid_indices:
                valid_indices.append(new_idx)
        
        # 获取定日镜实际位置
        mirror_positions = [self.candidate_positions[idx] for idx in valid_indices]
        tower_position = (tower_x, tower_y)
        
        return mirror_positions, tower_position
    
    def fitness_function(self, individual):
        """适应度函数 (目标函数) - 标准版本
        严格按照公式: max Σᴺᵢ₌₁(ηtotal,i × Amirror × DNIavg) / 总场地面积
        """
        try:
            # 解码个体
            mirror_positions, tower_position = self.decode_individual(individual)
            
            # 标准: 硬约束检查 - 不满足直接返回0
            if not self.check_constraints(mirror_positions, tower_position):
                return 0.001  # 不满足约束的个体适应度极低
            
            # 动态选择计算精度 - 最后20%代数使用完整计算
            use_full_efficiency = hasattr(self, 'current_generation') and \
                                self.current_generation > self.max_generations * 0.8
            
            if use_full_efficiency:
                # 使用完整效率计算 (精确但较慢)
                efficiencies = self._calculate_full_efficiency(
                    mirror_positions, tower_position
                )
            else:
                # 计算光学效率 - 使用简化计算提高速度
                efficiencies = self._calculate_simplified_efficiency(
                    mirror_positions, tower_position
                )
            
            # 标准目标函数: Σᴺᵢ₌₁(ηtotal,i × Amirror × DNIavg) / 总场地面积
            mirror_area = 36  # 6m × 6m = 36m²
            total_annual_power = sum(eta * mirror_area * self.dni_avg for eta in efficiencies)
            
            # 单位面积年平均输出功率 (标准目标函数)
            fitness = total_annual_power / self.field_area if self.field_area > 0 else 0
            
            return max(0.001, fitness)  # 确保适应度不为0
            
        except Exception as e:
            print(f"适应度计算错误: {e}")
            return 0.001  # 返回最小正值而不是0

    def _calculate_constraint_penalty(self, mirror_positions, tower_position):
        """计算约束违规惩罚 - 软约束策略"""
        penalty = 0.0
        tx, ty = tower_position
        
        # 1. 吸收塔位置约束
        if tx**2 + ty**2 > self.R_field**2:
            penalty += 0.5  # 50%惩罚
        
        # 2. 安全距离约束 - 软约束
        violation_count = 0
        total_pairs = 0
        
        for i, pos1 in enumerate(mirror_positions):
            for j, pos2 in enumerate(mirror_positions):
                if i >= j:
                    continue
                
                total_pairs += 1
                x1, y1 = pos1
                x2, y2 = pos2
                distance = math.sqrt((x1-x2)**2 + (y1-y2)**2)
                
                if distance < self.safety_distance:
                    violation_count += 1
        
        # 安全距离违规惩罚 (按比例)
        if total_pairs > 0:
            violation_ratio = violation_count / total_pairs
            penalty += min(0.3, violation_ratio * 0.5)  # 最多30%惩罚
        
        return min(0.8, penalty)  # 最大惩罚80%

    def _calculate_simplified_efficiency(self, mirror_positions, tower_position):
        """简化的效率计算 - 提高遗传算法速度"""
        efficiencies = []
        
        for mirror_pos in mirror_positions:
            # 简化计算: 只考虑几何因子和距离衰减
            mx, my = mirror_pos
            tx, ty = tower_position
            
            # 距离计算
            distance = math.sqrt((mx - tx)**2 + (my - ty)**2)
            
            # 简化效率模型 (基于经验公式)
            # 考虑距离衰减、角度因子的近似
            distance_factor = max(0.1, 1.0 / (1 + distance / 100))  # 距离衰减
            geometry_factor = max(0.3, 1.0 - distance / 500)  # 几何因子
            
            # 综合效率 (简化)
            efficiency = distance_factor * geometry_factor * 0.7  # 基准效率0.7
            efficiencies.append(max(0.05, min(0.9, efficiency)))
        
        return efficiencies
    
    def _calculate_full_efficiency(self, mirror_positions, tower_position):
        """完整效率计算 - 用于遗传算法后期精确评估"""
        # 使用效率计算器进行完整计算
        try:
            # 简化的太阳位置 (使用正确的字典格式)
            typical_solar_positions = [
                {'elevation': 45, 'azimuth': 180},  # 典型太阳高度角和方位角
                {'elevation': 60, 'azimuth': 180},
                {'elevation': 30, 'azimuth': 180}
            ]
            
            total_efficiencies = self.efficiency_calc.calculate_total_efficiency(
                mirror_positions, tower_position, typical_solar_positions
            )
            
            # 返回平均效率
            return [sum(eff_list)/len(eff_list) if eff_list else 0.1 
                   for eff_list in total_efficiencies]
            
        except Exception as e:
            print(f"完整效率计算失败，回退到简化计算: {e}")
            return self._calculate_simplified_efficiency(mirror_positions, tower_position)

    def check_constraints(self, mirror_positions, tower_position):
        """约束检查 - 标准版本，严格约束"""
        tx, ty = tower_position
        
        # 1. 吸收塔位置约束 (必须在场地内)
        if tx**2 + ty**2 > self.R_field**2:
            return False
        
        # 2. 定日镜数量约束 (标准: 100 ≤ N ≤ 2000)
        if len(mirror_positions) < self.min_mirrors or len(mirror_positions) > self.max_mirrors:
            return False
        
        # 3. 安全距离约束 - 标准: 严格检查所有定日镜对
        for i, pos1 in enumerate(mirror_positions):
            for j, pos2 in enumerate(mirror_positions):
                if i >= j:
                    continue
                
                x1, y1 = pos1
                x2, y2 = pos2
                distance = math.sqrt((x1-x2)**2 + (y1-y2)**2)
                
                # 标准: 任何违反安全距离的都不允许
                if distance < self.safety_distance:
                    return False
        
        # 4. 定日镜位置约束 (必须在场地内)
        for mx, my in mirror_positions:
            if mx**2 + my**2 > self.R_field**2:
                return False
        
        return True
    
    def tournament_selection(self, population, fitnesses, tournament_size=3):
        """锦标赛选择"""
        selected = []
        
        for _ in range(len(population)):
            # 随机选择tournament_size个个体
            tournament_indices = random.sample(range(len(population)), tournament_size)
            tournament_fitnesses = [fitnesses[i] for i in tournament_indices]
            
            # 选择适应度最高的个体
            winner_idx = tournament_indices[tournament_fitnesses.index(max(tournament_fitnesses))]
            selected.append(population[winner_idx].copy())
        
        return selected
    
    def crossover(self, parent1, parent2):
        """交叉操作 - 修复版本"""
        if random.random() > self.crossover_rate:
            return parent1.copy(), parent2.copy()
        
        # 解码父代
        N1, N2 = parent1[0], parent2[0]
        
        # 修复: 选择合适的定日镜数量
        max_possible = len(self.candidate_positions)
        N_child = min(min(N1, N2), max_possible)
        N_child = max(1, N_child)  # 确保至少有1个定日镜
        
        # 交叉定日镜位置索引
        indices1 = parent1[1:N1+1]
        indices2 = parent2[1:N2+1]
        
        # 修复: 处理索引长度不足的情况
        if len(indices1) == 0 or len(indices2) == 0:
            # 如果其中一个父代没有定日镜，直接返回原个体
            return parent1.copy(), parent2.copy()
        
        # 单点交叉
        crossover_point = random.randint(1, min(len(indices1), len(indices2)) - 1)
        
        child1_indices = indices1[:crossover_point] + indices2[crossover_point:crossover_point + (N_child - crossover_point)]
        child2_indices = indices2[:crossover_point] + indices1[crossover_point:crossover_point + (N_child - crossover_point)]
        
        # 去重处理
        child1_indices = list(dict.fromkeys(child1_indices))[:N_child]
        child2_indices = list(dict.fromkeys(child2_indices))[:N_child]
        
        # 补充不足的索引
        available_indices = list(range(len(self.candidate_positions)))
        
        while len(child1_indices) < N_child:
            new_idx = random.choice(available_indices)
            if new_idx not in child1_indices:
                child1_indices.append(new_idx)
        
        while len(child2_indices) < N_child:
            new_idx = random.choice(available_indices)
            if new_idx not in child2_indices:
                child2_indices.append(new_idx)
        
        # 交叉吸收塔位置
        tower1_x, tower1_y = parent1[N1+1], parent1[N1+2]
        tower2_x, tower2_y = parent2[N2+1], parent2[N2+2]
        
        alpha = random.random()
        child1_tower_x = alpha * tower1_x + (1 - alpha) * tower2_x
        child1_tower_y = alpha * tower1_y + (1 - alpha) * tower2_y
        
        child2_tower_x = (1 - alpha) * tower1_x + alpha * tower2_x
        child2_tower_y = (1 - alpha) * tower1_y + alpha * tower2_y
        
        # 构造子代染色体
        child1 = [N_child] + child1_indices + [child1_tower_x, child1_tower_y]
        child2 = [N_child] + child2_indices + [child2_tower_x, child2_tower_y]
        
        return child1, child2
    
    def mutation(self, individual):
        """变异操作 - 修复版本"""
        if random.random() > self.mutation_rate:
            return individual
        
        mutated = individual.copy()
        N = mutated[0]
        
        mutation_type = random.choice(['mirror_count', 'mirror_position', 'tower_position'])
        
        if mutation_type == 'mirror_count':
            # 修复: 动态调整定日镜数量范围
            max_possible = len(self.candidate_positions)
            actual_min = min(self.min_mirrors, max_possible)
            actual_max = min(self.max_mirrors, max_possible)
            
            if actual_min > actual_max:
                actual_min = max(1, actual_max // 2)
            
            # 变异定日镜数量
            new_N = random.randint(actual_min, actual_max)
            if new_N != N:
                if new_N > N:
                    # 增加定日镜
                    current_indices = mutated[1:N+1]
                    available_indices = [i for i in range(len(self.candidate_positions)) if i not in current_indices]
                    
                    # 修复: 确保有足够的可用索引
                    additional_needed = new_N - N
                    if len(available_indices) >= additional_needed:
                        additional_indices = random.sample(available_indices, additional_needed)
                        mutated = [new_N] + current_indices + additional_indices + mutated[N+1:]
                    else:
                        # 如果可用索引不够，只添加现有的
                        mutated = [N + len(available_indices)] + current_indices + available_indices + mutated[N+1:]
                else:
                    # 减少定日镜
                    current_indices = mutated[1:N+1]
                    if len(current_indices) >= new_N:
                        selected_indices = random.sample(current_indices, new_N)
                        mutated = [new_N] + selected_indices + mutated[N+1:]
        
        elif mutation_type == 'mirror_position':
            # 变异定日镜位置
            if N > 0:
                pos_to_change = random.randint(1, N)
                current_indices = mutated[1:N+1]
                available_indices = [i for i in range(len(self.candidate_positions)) if i not in current_indices]
                
                if available_indices:
                    mutated[pos_to_change] = random.choice(available_indices)
        
        elif mutation_type == 'tower_position':
            # 变异吸收塔位置
            tower_x = mutated[N+1] + random.gauss(0, 20)  # 高斯扰动
            tower_y = mutated[N+2] + random.gauss(0, 20)
            
            # 确保吸收塔在场地内
            if tower_x**2 + tower_y**2 <= self.R_field**2:
                mutated[N+1] = tower_x
                mutated[N+2] = tower_y
        
        return mutated
    
    def optimize(self):
        """执行遗传算法优化 - 修复版本"""
        print("\n🧬 开始遗传算法优化...")
        
        # 初始化种群
        population = self.initialize_population()
        
        best_fitness_history = []
        avg_fitness_history = []
        no_improvement_count = 0
        
        for generation in range(self.max_generations):
            self.current_generation = generation # 记录当前代数
            # 计算适应度
            fitnesses = []
            for individual in population:
                fitness = self.fitness_function(individual)
                fitnesses.append(fitness)
            
            # 统计信息
            best_fitness = max(fitnesses)
            avg_fitness = sum(fitnesses) / len(fitnesses)
            best_fitness_history.append(best_fitness)
            avg_fitness_history.append(avg_fitness)
            
            # 每10代输出进度
            if generation % 10 == 0 or generation == self.max_generations - 1:
                avg_fitness = sum(fitnesses) / len(fitnesses)
                computation_mode = "精确计算" if hasattr(self, 'current_generation') and \
                                          self.current_generation > self.max_generations * 0.8 else "简化计算"
                print(f"   第 {generation+1:3d} 代: 最佳={best_fitness:.6f}, 平均={avg_fitness:.6f} [{computation_mode}]")
            
            # 修复: 改进早停条件
            if generation > 0:
                improvement = best_fitness - best_fitness_history[-2]
                if improvement < 0.001:  # 改进阈值
                    no_improvement_count += 1
                else:
                    no_improvement_count = 0
                
                # 早停条件：连续20代无显著改善
                if no_improvement_count >= 20:
                    print(f"   早停: 连续{no_improvement_count}代无显著改善")
                    break
            
            # 选择
            selected_population = self.tournament_selection(population, fitnesses)
            
            # 交叉和变异
            new_population = []
            for i in range(0, len(selected_population), 2):
                parent1 = selected_population[i]
                parent2 = selected_population[min(i+1, len(selected_population)-1)]
                
                child1, child2 = self.crossover(parent1, parent2)
                child1 = self.mutation(child1)
                child2 = self.mutation(child2)
                
                new_population.extend([child1, child2])
            
            population = new_population[:self.population_size]
        
        # 返回最佳解
        final_fitnesses = [self.fitness_function(ind) for ind in population]
        best_index = final_fitnesses.index(max(final_fitnesses))
        best_solution = population[best_index]
        
        print(f"✅ 遗传算法优化完成")
        print(f"   最终最佳适应度: {max(final_fitnesses):.4f}")
        print(f"   总迭代次数: {generation + 1}")
        
        return best_solution, max(final_fitnesses)

# 测试遗传算法模块
if __name__ == "__main__":
    print("🚀 开始遗传算法测试...")
    
    # 创建测试数据
    candidate_positions = [(i*10, j*10) for i in range(-30, 31) for j in range(-30, 31) 
                          if 100 <= i*10**2 + j*10**2 <= 350**2]
    
    solar_positions = [
        {'elevation': math.radians(45), 'azimuth': math.radians(180)},
        {'elevation': math.radians(60), 'azimuth': math.radians(200)}
    ]
    
    dni_avg = 600  # W/m²
    
    # 初始化遗传算法 (小规模测试)
    ga = GeneticAlgorithm(candidate_positions[:100], solar_positions, dni_avg)
    ga.population_size = 10
    ga.max_generations = 5
    
    # 运行优化
    best_solution, best_fitness, fitness_history, avg_history = ga.evolve()
    
    print(f"✅ 测试完成，最优解适应度: {best_fitness:.4f}") 